[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel - #3315
Open
cael-ling wants to merge 5 commits into
Open
Conversation
Add nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with columnwise MXFP8 quantization over grouped (MoE) tensors. Signed-off-by: Cael Ling <caell@nvidia.com>
cael-ling
requested review from
Oleg-Goncharov,
ksivaman and
ptrendx
as code owners
August 4, 2026 14:36
github-actions
Bot
added
the
community-contribution
for more information, see https://pre-commit.ci
Contributor
Greptile SummaryThe PR adds a grouped weighted-SwiGLU operation fused with columnwise MXFP8 quantization for MoE FC2 weight-gradient inputs.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains; the previously reported stride, input-device, and metadata-device issues are guarded in the new binding, while the prior current-device metadata rejection no longer exists in the shared validator. Important Files Changed
Sequence DiagramsequenceDiagram
participant Python as PyTorch caller
participant Binding as group_swiglu_quantize binding
participant Core as NVTE dispatch
participant Kernel as MXFP8 CUDA kernel
Python->>Binding: input_2f, prob, quantizer, metadata
Binding->>Binding: Validate layout, dtype, and device
Binding->>Binding: Allocate grouped [T,F] output
Binding->>Core: nvte_group_swiglu_quantize(...)
Core->>Kernel: Dispatch weighted SwiGLU + columnwise quantization
Kernel-->>Python: Grouped columnwise MXFP8 output
Reviews (4): Last reviewed commit: "[PyTorch] Scope the grouped metadata dev..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
…tize Signed-off-by: Cael Ling <caell@nvidia.com>
Signed-off-by: Cael Ling <caell@nvidia.com>
…ntize Signed-off-by: Cael Ling <caell@nvidia.com>
vthumbe1503
reviewed
Aug 7, 2026
vthumbe1503
left a comment
vthumbe1503
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the PR @cael-ling. Could you also please publish the perf numbers for this kernel?
| * \param[in,out] output Grouped output tensor [T, F] (columnwise MXFP8). | ||
| * \param[in] stream CUDA stream used for the operation. | ||
| */ | ||
| void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob, |
Collaborator
There was a problem hiding this comment.
Can we rename this? Quantization details are already encoded in GroupedTensor. It ll be consistent with the other API naming convention
Suggested change
| void nvte_group_swiglu_quantize(const NVTEGroupedTensor input, const NVTETensor prob, | |
| void nvte_group_scaled_swiglu(const NVTEGroupedTensor input, const NVTETensor prob, |
| std::optional<at::Tensor> tensor_offsets, | ||
| std::optional<at::Tensor> noop_flag); | ||
|
|
||
| py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob, |
Collaborator
There was a problem hiding this comment.
Similar comment.
Suggested change
| py::object group_swiglu_quantize(const at::Tensor &input_2f, const at::Tensor &prob, | |
| py::object group_scaled_swiglu(const at::Tensor &input_2f, const at::Tensor &prob, |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add
nvte_group_swiglu_quantize, which fuses the weighted SwiGLU activation with column-wise MXFP8 quantization over grouped (MoE) tensors:Type of change
Changes
Please list the changes introduced in this PR:
Checklist: